import networkx as nx
import numpy as np
import glob
import os, os.path
import math
import pickle
import scipy.sparse as sp
import random
from sklearn.decomposition import PCA

import Link_Prediction_Scores_All
#line  论文中  无向图结合使用一阶相似度和二阶相似度
#一阶相似度  相近节点的欧式距离尽可能小
#二阶相似度  使得自编码器邻接矩阵和编码后的向量尽可能相似

#模拟的相似度矩阵 一阶用邻接矩阵  二阶用common neighbors或 Adamic-Adar

DEAL_NAME = ['Amherst','Facebook','hamster','Oberlin','PB','SC_TS','USAir']

def getOneOrderProbabilityMatrix(adj):
   adj_ = adj.toarray()
   return adj_

def getSecondOrderProbabilityMatrix_AA(adj):
   num_nodes = adj.shape[0]
   adj_ = adj.toarray()
   G = nx.from_numpy_matrix(adj_)
   piter = nx.adamic_adar_index(G)
   S = np.zeros((num_nodes,num_nodes))
   for u, v, p in piter:
      S[u][v] = p
      S[v][u] = p
   for i in range(num_nodes):
       degree = G.degree(i)
       row_sum = np.sum(S[i])
       for j in range(num_nodes):
          if(S[i][j]!=0):
             S[i][j] = degree*(S[i][j]/row_sum)
   return S

def getSecondOrderProbabilityMatrix_CN(adj):
    S = adj.dot(adj)
    return S.toarray()

def getProbabilityMatrix(adj):
   S = getOneOrderProbabilityMatrix(adj)+getSecondOrderProbabilityMatrix_AA(adj)
   return S

def pca_probabilityMatrix(M,components):
   pca = PCA(n_components=components)
   pca.fit(M)
   newM = pca.fit_transform(M)
   variance_ratio_ = pca.explained_variance_ratio_
   print(pca.explained_variance_ratio_)
   print(np.sum(variance_ratio_))
   return newM

prefix_name = 'D:/data/similarity_matrix/Line/'



#保存相似矩阵和pca后的
def save_Smatrix_newS(adj,dela_name,components):
   ori_Smatirx_name = dela_name + '-line-ori.pkl'
   pca__Smatirx_name = dela_name + '-line-pca-{}.pkl'.format(components)
   file_name = prefix_name + ori_Smatirx_name
   if(not os.path.exists(file_name)):
      S = getProbabilityMatrix(adj)
      with open(file_name, 'wb') as f:
        pickle.dump(S, f, protocol=2)
      print(file_name+"S orimatrix saved!")
   else:
      with open(file_name, 'rb') as f:
         S = pickle.load(f)
   newS = pca_probabilityMatrix(S, components)
   file_name = prefix_name + pca__Smatirx_name
   with open(file_name, 'wb') as f:
      pickle.dump(newS, f, protocol=2)
   print(file_name+"S pcamatrix saved!")


if __name__ == '__main__':
   for dela_name in DEAL_NAME:
      combined_dir = 'D:/data-processed/{}-adj.pkl'.format(dela_name)
      with open(combined_dir, 'rb') as f:
         adj = pickle.load(f)
         S = getSecondOrderProbabilityMatrix_CN(adj)
         if(adj.shape[0]<=1000):
            save_Smatrix_newS(adj, dela_name, 48)
         elif(adj.shape[0]<=4000):
            save_Smatrix_newS(adj, dela_name, 96)
         elif(adj.shape[0]<=10000):
            save_Smatrix_newS(adj,dela_name,192)
   # pca_Smatirx_name = 'Oberlin-pca.pkl'
   # with open(file_name, 'rb') as f:
   #    S = pickle.load(f)
   # file_name = prefix_name + pca_Smatirx_name
   # with open(file_name, 'rb') as f:
   #    newS = pickle.load(f)
   # print('')



